import yolov3_model
from dataset import *
from loss import *
import paddle
from paddle.optimizer import Adam
from paddle.nn.functional import binary_cross_entropy_with_logits

batch_size = 4
numclasses = 4
LR = 1e-4
Epoch = 100

anchors_path = ''
dataPath = ''
inputShape = (416, 416)

anchor = [[1, 1], [1, 3], [3, 2], [3, 6], [6, 4], [6, 12], [12, 9], [16, 21], [40, 35]]


dataset = myDataSet(r'data', anchor)
dataloader = DataLoader(dataset, batch_size=batch_size, shuffle=True)

model = yolov3_model.yolov3()
scheduler = paddle.optimizer.lr.NoamDecay(
    d_model=0.01, warmup_steps=100, verbose=True)

optimizer = Adam(parameters=model.parameters(), learning_rate=LR)
#device = paddle.device("cuda:0"if paddle.cuda.is_available() else "cpu")
paddle.device.set_device('cpu')
model.train()

def train(epoch):
    all_loss = 0.0
    for index, data in enumerate(dataloader, 0):
        img = data[0]
        img = img.astype(paddle.float32)
        out = model(img)
        y_true = (data[1], data[2], data[3])

        myloss = loss(anchor)
        lastloss = myloss(out, y_true)

        lastloss.backward()
        optimizer.step()
        #optimizer.minimize(lastloss)
        optimizer.clear_grad()
        all_loss += lastloss.item()
        if index % 5 == 0:
            print('[epoch:%3d];index:%3d;loss=%.4f'%(epoch+1, index, all_loss))
    if (epoch+1) % 10 == 0:
        paddle.save(model.state_dict(), path='weight.pdparams')

for epoch in range(Epoch):
    train(epoch)











